import cv2
import glob
import os
import os.path as osp
import mmcv
import sys

def showImg(img,wati_time=0,win_name="r"):
    print(img.shape)
    cv2.namedWindow(win_name,cv2.WINDOW_NORMAL)
    # cv2.resizeWindow(win_name,1920,1080)
    cv2.imshow(win_name,img)
    key = cv2.waitKey(wati_time) & 0xFF
    if key == ord('q'):
        sys.exit()
    return key

def yaml_naive(w,h,K,D, save_dir='.'):

    fs = cv2.FileStorage(os.path.join(save_dir,'KD.yaml'), cv2.FILE_STORAGE_WRITE)
    fs.write('Width', w)
    fs.write('Height', h)
    fs.write('K',K)
    fs.write('D',D)
    fs.release()


root1280 = "/cv/zjh/well_driver/log/ws/data_record_sync/eval_bev/qtruck21/K"
srcdata = "/cv/zjh/well_driver/log/ws/data_record_sync/eval_bev/qtruck21/data"
Ks1280  = glob.glob(osp.join(root1280,"*yaml"))
Ks1280.sort()
srcdata_l  = glob.glob(osp.join(srcdata,"*jpg"))
srcdata_l.sort()


for i,K_p in enumerate(Ks1280):
    fs  = cv2.FileStorage(K_p, cv2.FileStorage_READ)
    K = fs.getNode('K').mat()
    D = fs.getNode('D').mat()
    print(srcdata_l[i])
    src = cv2.imread(srcdata_l[i])
    undistorted_img = cv2.undistort(src, K, D)   
    showImg(undistorted_img,0,"rr")
    
    K = K*3/2
    K[2][2]=1.0
    mmcv.mkdir_or_exist(K_p[:-5])
    w= 1920
    h=1080
    print(K_p[:-5])
    yaml_naive(w,h,K,D,K_p[:-5])
    print(K_p[:-5]+"/KD.yaml")
    
    

